from pyaudio import PyAudio, paInt16, paContinue, paComplete  
import wave 
import time
import numpy as np 

chunk = 1024
format_sample = paInt16 
channels = 1
fps = 16000 
# file to record
wf = wave.open("./test.wav", "wb")

# a ring buffer
class ring_buffer():
    def __init__(self, buffer, size):
        self.buffer = buffer
        self.size = size
        self.rNdx = 0
        self.wNdx = 0
        self.cnt = 0
    
    def GetFreeBytes(self):
        return self.size - self.cnt 

    def GetUsedBytes(self):
        return self.cnt 
    
    def Write(self, data_w, dataBytes):
        writeToEnd = self.size - self.wNdx
        bytesToCopy = min(dataBytes, self.size - self.cnt)
        if(bytesToCopy > 0):
            CopyToEnd = min(bytesToCopy, writeToEnd)
            self.buffer[self.wNdx : self.wNdx + CopyToEnd] = data_w[:CopyToEnd]
            if(writeToEnd < bytesToCopy):
                self.buffer[:(bytesToCopy - writeToEnd)] = data_w[writeToEnd:]
            # update the wNdx 
            self.wNdx = (self.wNdx + bytesToCopy) % self.size
            self.cnt = (self.cnt + dataBytes) #% self.size
            
        return bytesToCopy 
    
    def read(self, data_r, dataBytes, deleteBytes = 0):
        readToEnd = self.size - self.rNdx 
        bytesToCopy = min(dataBytes, self.cnt)
        if(bytesToCopy != 0):
            CopyToEnd = min(dataBytes, readToEnd)
            data_r[:CopyToEnd] = self.buffer[self.rNdx:self.rNdx+CopyToEnd]
            if(readToEnd < bytesToCopy):
                data_r[readToEnd:] = self.buffer[:bytesToCopy - readToEnd]
            rNdx_offset = bytesToCopy if deleteBytes == 0 else deleteBytes
            self.rNdx = (self.rNdx + rNdx_offset) % self.size
            self.cnt -= rNdx_offset
        
        return bytesToCopy
 
buffer = np.zeros(fps * channels + 1024 * 10 ,dtype='int16')
ring = ring_buffer(buffer, len(buffer))
               
def callback(in_data, frame_count, time_info, status):
    #wf.writeframes(in_data)
    ring.Write(np.frombuffer(in_data, dtype='int16'), len(in_data) // 2) # in_data is bytesarray
    return in_data, paContinue

p = PyAudio() 
stream = p.open(
                format=format_sample, channels=channels,
                rate=fps,
                frames_per_buffer=chunk,
                input=True,
                output=True,
                input_device_index = None,
                stream_callback = callback 
                )

# declare the wave handler
wf.setnchannels(channels)
wf.setsampwidth(p.get_sample_size(format_sample))
wf.setframerate(fps)

if __name__ == "__main__":
    audio_data_1s = np.zeros(fps * channels, dtype='int16')
    audio_idx = 0

    audio_data = b''
    stream.start_stream() 
    while(stream.is_active()):
        # audio_in = stream.read(chunk)
        # audio_data += audio_in
        # stream.write(audio_in)
        if(ring.GetUsedBytes() >= fps * channels):
            ring.read(audio_data_1s, len(audio_data_1s))
            with wave.open("./wave_files/%d.wav"%audio_idx, "wb") as wf:
                audio_idx += 1
                wf.setnchannels(channels)
                wf.setsampwidth(p.get_sample_size(format_sample))
                wf.setframerate(fps)
                wf.writeframes(audio_data_1s)
                wf.close()
            
        #     print()
    stream.stop_stream()
    stream.close()
    wf.close()
    p.terminate()
